import pysac.mhs_atmosphere as atm

#Read in the VAL3C model
empirical_data = atm.hs_atmosphere.read_VAL3c_MTW(MTW_file=False)


# Create a Z array at the interpolated resolution and interpolate.
ZZ = u.Quantity(np.linspace(empirical_data['Z'][0], empirical_data['Z'][-1], 128), unit=empirical_data['Z'].unit)
table = atm.hs_atmosphere.interpolate_atmosphere(empirical_data, ZZ, s=0)


# Create a figure and make space for the axes.
fig, ax = plt.subplots(gridspec_kw={'right':0.77, 'left':0.16, 'bottom':0.13}, figsize=(13,8))

# Shortcut all the Mm conversion.
Z = empirical_data['Z'].to(u.Mm)

lrho, = ax.plot(Z, empirical_data['rho'].quantity.si, 'x', color='blue')
lrho_i, = ax.plot(ZZ.to(u.Mm), table['rho'].quantity.si, color='blue')

ax2 = ax.twinx()
lp, = ax2.plot(Z, empirical_data['p'].to(u.Pa), 'x', color='green')
lp_i, = ax2.plot(ZZ.to(u.Mm), table['p'].to(u.Pa), color='green')


ax3 = ax.twinx()
ax3.spines["right"].set_position(("axes", 1.2))
lt, = ax3.plot(Z, empirical_data['T'].to(u.K), 'x', color='red')
lt_i, = ax3.plot(ZZ.to(u.Mm), table['T'].to(u.K), color='red')


# Set primary axes properties and labels
ax.semilogy()
ax.set_ylabel(r"Density [{}]".format(lrho._yorig.unit._repr_latex_()))
ax.set_xlabel(r"Height [{}]".format(lrho._xorig.unit._repr_latex_()))
ax.set_xlim(Z[0].value-0.1, Z[-1].value+0.1)


# Pressure Axis
ax2.semilogy()
ax2.set_ylabel(r"Pressure [{}]".format(lp._yorig.unit._repr_latex_()))


# Temp axis
ax3.set_ylabel(r"Temperature [{}]".format(lt._yorig.unit._repr_latex_()))

ax.set_xlim([-0.02,1.62])
ax3.set_ylim([3500,7500])
# Set the colours for the ticks and the labels.
ax.tick_params(axis='y', colors=lrho.get_color())
ax2.tick_params(axis='y', colors=lp.get_color())
ax3.tick_params(axis='y', colors=lt.get_color())

ax.yaxis.label.set_color(lrho.get_color())
ax2.yaxis.label.set_color(lp.get_color())
ax3.yaxis.label.set_color(lt.get_color())
fig

slc = yt.SlicePlot(ds, 'x', 'density', origin='lower-center-domain', axes_unit='Mm')
slc.set_figure_size(10)
slc.set_cmap('all', 'viridis')
slc.set_font_size(20)
slc.set_zlim('all', 0, 2.5e-4)


seed_points = np.zeros([11,2]) + 1.52
seed_points[:,0] = np.linspace(-0.99, 0.95, seed_points.shape[0], endpoint=True)
slc.annotate_streamlines('mag_field_y', 'mag_field_z', field_color='magnetic_field_strength',
plot_args={'start_points':seed_points, 'density':15, 'cmap':'Blues', 'linewidth':2,
'norm': matplotlib.colors.LogNorm(*ds.all_data().quantities.extrema("magnetic_field_strength"))})
#force render
slc.save('/tmp/test.png')
#use the raw Figure for transparent bg
slc.plots['density'].figure

xc = domain['xmax']/2
yc = domain['ymax']/2
ti = 0
n = 100

surf_seeds = []
for theta in np.linspace(0, 2 * np.pi, n, endpoint=False):
    surf_seeds.append([tube_r * np.cos(theta + 0.5 * ti) + xc,
    tube_r * np.sin(theta + 0.5 * ti) + yc, domain['zmax']])

seeds = np.array(surf_seeds)
#Add axes:
axes, outline = mpf.add_axes(np.array(zip(ds.domain_left_edge,ds.domain_right_edge)).flatten()/1e8)

#Add seed points to plot:
seed_points = mlab.points3d(seeds[:,0], seeds[:,1], seeds[:,2],
color=(0.231, 0.298, 0.752), scale_mode='none',
scale_factor=1.5)

mlab_view(fig.scene)

field_lines = SeedStreamline(seed_points = np.array(seeds))
bmag.add_child(field_lines)
field_lines.actor.mapper.scalar_visibility = False
field_lines.actor.property.color = (0,0,0)
field_lines.actor.property.line_width = 1.5

mlab_view(fig.scene)

pd_seeds = ttf.make_circle_seeds(100, 60, **domain)
fieldlines, surface = ttf.create_flux_surface(bfield.outputs[0], pd_seeds)
surface.output.lines = None
flux_surface = mlab.pipeline.surface(surface.output)
flux_surface.actor.mapper.scalar_visibility = False
flux_surface.actor.property.color = (0.8,0.8,0.8)
#flux_surface.actor.property.line_width = 0

mlab_view(fig.scene)

axes.visible = False
outline.visible = False
flux_surface.actor.property.edge_visibility = True
mlab_view(fig.scene, azimuth = 90, elevation = 75, distance=80, focalpoint=[63, 120, 110], aa=20)

poly_norms = ttf.make_poly_norms(surface.output)
normvec = mlab.pipeline.glyph(poly_norms.output)
normvec.glyph.glyph_source.glyph_source = normvec.glyph.glyph_source.glyph_dict['arrow_source']
normvec.glyph.glyph.scale_mode = 'data_scaling_off'
normvec.glyph.glyph.color_mode = 'color_by_scale'
normvec.glyph.glyph.scale_factor = 5
normvec.glyph.glyph_source.glyph_position = 'tail'

mlab_view(fig.scene, azimuth=85, elevation=80, distance=50, focalpoint=[63, 120, 110], aa=20)

from streamlines import Streamlines
#Use Equation 1 to calculate the vector field in a 2D plane to plot it.
time = np.linspace(0,60,480)
dt = time[1:] - time [:-1]
period = 240.

x = np.linspace(7812.5,1992187.5,128)
y = np.linspace(7812.5,1992187.5,128)

x_max = x.max()
y_max = y.max()

xc = 1.0e6
yc = 1.0e6

xn = x - xc
yn = y - yc

delta_x=0.1e6
delta_y=0.1e6

xx, yy = np.meshgrid(xn,yn)
exp_y = np.exp(-(yn**2.0/delta_y**2.0))
exp_x = np.exp(-(xn**2.0/delta_x**2.0))

exp_x2, exp_y2= np.meshgrid(exp_x,exp_y)
exp_xyz = exp_x2 * exp_y2


#==============================================================================
# Define Driver Equations and Parameters
#==============================================================================
#A is the amplitude, B is the spiral expansion factor
A = 10

#Tdamp defines the damping of the driver with time, Tdep is the ocillator
tdamp = lambda time1: 1.0 #*np.exp(-(time1/(period)))
tdep = lambda time1: np.sin((time1*2.0*np.pi)/period) * tdamp(time1)

#Define a peak index to use for scaling in the inital frame
max_ind = np.argmax(tdep(time) > 0.9998)

def log():
    B = 0.05
    phi = np.arctan2(1,B)
    theta = np.arctan2(yy,xx)

    uy = np.sin(theta + phi)
    ux =  np.cos(theta + phi)

    vx = lambda time1: (ux / np.sqrt(ux**2 + uy**2)) * exp_xyz * tdep(time1) * A
    vy = lambda time1: (uy / np.sqrt(ux**2 + uy**2)) * exp_xyz * tdep(time1) * A

    vv = np.sqrt(vx(time[max_ind])**2 + vy(time[max_ind])**2)

    return vx, vy, vv

def arch():
    B = 0.005
    r = np.sqrt(xx**2 + yy**2)

    vx = lambda time1: ( (B*1e6 * xx) / (xx**2 + yy**2) + yy/r ) * exp_xyz * tdep(time1) * A
    vy = lambda time1: ( (B*1e6 * yy) / (xx**2 + yy**2) - xx/r ) * exp_xyz * tdep(time1) * A

    vv = np.sqrt(vx(time[max_ind])**2 + vy(time[max_ind])**2)

    return vx, vy, vv

def uniform():
    #Uniform
    vx = lambda time1: A * (yy / np.sqrt(xx**2 + yy**2)) * exp_xyz * tdep(time1)
    vy = lambda time1: A * (-xx / np.sqrt(xx**2 + yy**2)) * exp_xyz * tdep(time1)
    vv = np.sqrt(vx(time[max_ind])**2 + vy(time[max_ind])**2)
    
    return vx, vy, vv

drivers = [log, arch, uniform]

#fig, axs = plt.subplots(1, 3, figsize=(18,9))

for driver_func in drivers:
    fig, ax = plt.subplots(figsize=(7,6), dpi=300)
    #============================================================================
    # Do the Plotting
    #============================================================================
    vx, vy, vv = driver_func()
    # Calculate Streamline
    slines = Streamlines(x,y,vx(time[max_ind]),vy(time[max_ind]),maxLen=7000,
                         x0=xc, y0=yc, direction='forwards')

    im = ax.imshow(vv, cmap='Blues', extent=[7812.5,x_max,7812.5,y_max])
    im.set_norm(matplotlib.colors.Normalize(vmin=0,vmax=A))
    #ax.hold()
    
    if driver_func != uniform:
        Sline, = ax.plot(slines.streamlines[0][0],slines.streamlines[0][1],color='red',linewidth=2, zorder=40)
    else:
        Sline = matplotlib.patches.Circle([1e6, 1e6], radius=.15e6, fill=False, color='red', linewidth=2, zorder=40)
        ax.add_artist(Sline)

    #Add colourbar
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.2)
    cbar = plt.colorbar(im,cax)
    cbar.set_label(r"$|V|$ [ms$^{-1}$]")
    scalar = matplotlib.ticker.ScalarFormatter(useMathText=False,useOffset=False)
    scalar.set_powerlimits((-3,3))
    cbar.formatter = scalar
    cbar.ax.yaxis.get_offset_text().set_visible(True)
    cbar.update_ticks()
    #cbar.solids.set_rasterized(True)
    cbar.solids.set_edgecolor("face")

    #Add quiver plot overlay
    qu = ax.quiver(x, y, vx(time[max_ind]), vy(time[max_ind]), scale=25*A, color='k', zorder=20, linewidth=1)
    ax.axis([8.0e5,12.0e5,8.0e5,12.0e5])

    ax.xaxis.set_major_formatter(scalar)
    ax.yaxis.set_major_formatter(scalar)
    ax.xaxis.set_major_locator(matplotlib.ticker.MaxNLocator(5))
    ax.yaxis.set_major_locator(matplotlib.ticker.MaxNLocator(5))
    ax.xaxis.get_offset_text().set_visible(False)
    ax.yaxis.get_offset_text().set_visible(False)
    ax.set_xlabel("X [Mm]")
    ax.set_ylabel("Y [Mm]")

    fig.savefig('images/driver_{}.png'.format(driver_func.__name__), transparent=True)

figures['vert']

figures['horiz']

figures['Slog']

from flux_comparison import make_flux_bar_chart, get_averages

Favgs = get_averages(ch4.data_dir)

fig = make_flux_bar_chart((13,11), ch4.data_dir)
fig

fig, ax = plt.subplots(figsize=(14,2), dpi=600)
ax.plot(BL, np.ones(BL.size), 'x', markersize=10, mew=2)
ax.errorbar([0.15], [1], xerr=np.array([[-1*(0.15-1/(6.4-1.6)), 0.15+1/(6.4+1.6)]]).T, mew=2, elinewidth=2)
ax.semilogx()
ax.get_yaxis().set_visible(False)
ax.set_frame_on(False)
ax.get_xaxis().tick_bottom()
ax.xaxis.set_tick_params(width=2)
ax.xaxis.set_tick_params(width=2, which='minor')
ax.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
ax.xaxis.set_ticks(BL)
xmin, xmax, ymin, ymax = ax.axis()
ax.add_artist(plt.Line2D((xmin, xmax), (ymin, ymin), color='black', linewidth=1.4))
l = ax.set_xlim([0.01, 2.0])
l = ax.set_xlabel(r'$B_L$', fontsize=18)

fig.tight_layout(h_pad=0.01)
fig

bl_figures['B0015']

bl_figures['B015']

bl_figures['B15']

import os
int_periods = np.floor(600./cfg.period)*180

def calc_avgs(tube_r):
    AvgsP = np.zeros([3,len(BL)])
    for i, bl in enumerate(BL):
        cfg.exp_fac = bl
        
        times = np.load(os.path.join(cfg.data_dir, 'Times_{}.npy'.format(cfg.get_identifier())))
        max_index = np.argmin(np.abs(int_periods - times))
        
        Fpar, Fperp, Fphi = map(np.load, ph.glob_files(cfg, tube_r, 'LineFlux*Fp*npy'))
        #Fpar, Fperp, Fphi = map(np.load, ph.glob_files(cfg, tube_r, '*vp*npy'))
        Fpar[np.abs(Fpar)<1e-5], Fperp[np.abs(Fperp)<1e-5], Fphi[np.abs(Fphi)<1e-5] = 0., 0., 0.
        Fpar, Fperp, Fphi = Fpar[:max_index,:], Fperp[:max_index,:], Fphi[:max_index,:]
        
        Ftot2 = (Fpar**2 + Fperp**2 + Fphi**2)
        Fpar2, Fperp2, Fphi2 = np.array([Fpar, Fperp, Fphi])**2
        FparP, FperpP, FphiP = (Fpar2/Ftot2)*100, (Fperp2/Ftot2)*100, (Fphi2/Ftot2)*100
        
        FparP = np.mean(np.ma.masked_array(FparP, np.isnan(FparP)))
        FperpP = np.mean(np.ma.masked_array(FperpP, np.isnan(FperpP)))
        FphiP = np.mean(np.ma.masked_array(FphiP, np.isnan(FphiP)))
        
        AvgsP[:, i] = FparP, FperpP, FphiP
    return AvgsP

figsize = (9.5,11)
fig, axs = plt.subplots(nrows=3, figsize=figsize, sharex=True)
titles = ["Flux Surface Radius $=156$ km", "Flux Surface Radius $=468$ km", "Flux Surface Radius $=936$ km"]
tubes = []
for i, ax in enumerate(axs):
    AvgsP = calc_avgs(cfg.tube_radii[i])
    tubes.append(AvgsP)
    ax.semilogx()
    ax.plot(BL, AvgsP[0], 'o', label=r"$F_\parallel^2$", mew=0, ms=10)
    ax.plot(BL, AvgsP[1], '_', label=r"$F_\perp^2$", mew=2, ms=10)
    ax.plot(BL, AvgsP[2], 'x', label=r"$F_\phi^2$", mew=2, ms=10)
    ax.set_ylabel("% Square Wave Flux")
    ax.set_title(titles[i])
    ax.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
    ax.xaxis.set_ticks(BL)
    ax.set_xlim([0.01, 2.01])
    ax.set_ylim([0, 85])
    err = np.array([-1*(0.15-1/(6.4-1.6)), 0.15+1/(6.4+1.6)])
    ax.fill_betweenx(np.linspace(-5,105), err[0], err[1], alpha=0.3, color='green', linewidth=0)

axs[0].legend(loc=9)
axs[-1].set_xlabel("Logarithmic Spiral Expansion Factor ($B_L$)")

plt.tight_layout()
fig

beta = False
cfg.exp_fac = 0.15
ph.xxlim = 600
tube_r = 'r30'

def add_all_bpert(ax, tube_r, N=4, levels=None):
    kwargs = ph.get_triple(cfg, beta=beta, single='bpert')
    x = kwargs['x_{}'.format(tube_r)]
    y = kwargs['y_{}'.format(tube_r)]
    par = kwargs['par_line_{}'.format(tube_r)].T[::-1, :]
    par[np.abs(par)<=1e-12] = 0
    perp = kwargs['perp_line_{}'.format(tube_r)].T[::-1, :]
    perp[np.abs(perp)<=1e-12] = 0
    phi = kwargs['phi_line_{}'.format(tube_r)].T[::-1, :]
    phi[np.abs(phi)<=1e-12] = 0
    ax[0].contour(x, y, par, N, colors='k', linewidths=np.linspace(0.5,1.5,N))
    ax[1].contour(x, y, perp, N, colors='k', linewidths=np.linspace(0.5,1.5,N))
    ax[2].contour(x, y, phi, N, colors='k', linewidths=np.linspace(0.5,1.5,N))	                   

def add_triple_phase(ax, tube_r):
    ps = ph.get_phase_speeds(cfg, tube_r)
    for ax0 in ax:
        ph.add_phase_speeds(ax0, color='g', **ps)

captions = {p: r"Period: ${}$ s amplitude: ".format(p) + a + r" ms$^{{-1}}$" for p, a in zip(periods, str_amps)[:10]}
#print(captions, file=sys.stderr)
width = 0.79

p_figures= {}
for i, paf in enumerate(all_periods):
    [setattr(cfg, f, getattr(paf, f)) for f in paf._fields]

    fig, ax = plt.subplots(nrows=3, ncols=2, figsize=(17,8), sharex=True)
    
    kwargs = ph.get_single_velocity(cfg, 'r30', beta=beta)
    kwargs.update(pvel_labels)
    
    ph.triple_plot(ax[:,0], **kwargs)
    
    kwargs = ph.get_single_percentage_flux(cfg, 'r30', beta=beta)
    kwargs.update(pflux_labels)
    kwargs.update({'cmap': 'PRGn'})
    
    ph.triple_plot(ax[:,1], **kwargs)
    
    #Remove the top two x labels
    ax[0,0].set_xlabel('')
    ax[1,0].set_xlabel('')
    add_triple_phase(ax[:,0], 'r30')
    ax[0,1].set_xlabel('')
    ax[1,1].set_xlabel('')
    add_triple_phase(ax[:,1], 'r30')
    #add_all_bpert(ax, 'r30')
    fig.tight_layout(h_pad=0.05)
    p_figures[str(paf.period)] = fig

p_figures['30.0']

p_figures['150.0']

p_figures['300.0']

from period_amps import periods, sim_params
sim_params = sim_params[:10]
periods = np.array(periods[:10])

fig, axs = plt.subplots(nrows=3, figsize=(10,11), sharex=True)
titles = ["Flux Surface Radius $=156$ km", "Flux Surface Radius $=468$ km", "Flux Surface Radius $=936$ km"]
tubes = []
for i, ax in enumerate(axs):
    AvgsP = ph.get_all_avgs(cfg, cfg.tube_radii[i], sim_params)
    tubes.append(AvgsP)
    ax.plot(periods, AvgsP[0], 'o', label=r"$F_\parallel^2$", mew=0, ms=12)
    ax.plot(periods, AvgsP[1], '_', label=r"$F_\perp^2$", mew=2, ms=12)
    ax.plot(periods, AvgsP[2], 'x', label=r"$F_\phi^2$", mew=2, ms=12)
    ax.set_ylabel("% Square Wave Flux")
    ax.set_title(titles[i])
    ax.xaxis.set_major_formatter(matplotlib.ticker.ScalarFormatter())
    ax.xaxis.set_ticks(periods)
    ax.set_xticklabels(["{}\n[{:n}]".format(p, 600//p) for p in periods])
    ax.set_ylim([10, 75])
    ax.set_xlim([25, 305])

axs[-1].set_xlabel("Period [s] \n [Number of periods averaged]")
axs[0].legend(loc=2)

#axs[0].legend(bbox_to_anchor=(1.06, 1.05))
plt.tight_layout(h_pad=0.1)
fig